"""
FeatureGenerators.py - Provides methods to generate features for a review.

Author: Sophie Song
"""

import collections
import nltk
import operator
import random

from nltk.tag.simplify import simplify_wsj_tag

################################################################################
# FeatureGenerator                                                             #
# Base class, provides sanatize method.                                        #
################################################################################
class FeatureGenerator(object):
  def _sanitize(self, s):
    """Some characters will make Weka barf if they're in the arff file.
    This function replaces the bad characters with placeholders.

    Parameters:
      s - The string to sanitize

    Return:
      The string with the barf characters removed.
    """
    sub_dict = {
      '"': '<DQ>',
      "'": '<SQ>',
      ',': '<CO>',
      ' ': '<SP>',
      '\n': '<NL>',
      '@': '<AT>',
      '%': '<PR>',
      chr(1): '',
      chr(18): ''
    }
    x = ''.join(map(lambda x: sub_dict[x] if x in sub_dict else x, s))
    if type(x) == type(u''):
      return x
    else:
      return unicode(x, errors='ignore')


################################################################################
# UnigramFeatureGenerator                                                      #
# A simple feature generator that just considers unigrams.                     #
################################################################################
class UnigramFeatureGenerator(FeatureGenerator):
  def __init__(self, data_fetcher, cutoff=20):
    """Constructor.

    Parameters:
      data_fetcher - The data_fetcher for the corpus this feature generator will
                     be used with.
      cutoff       - The minimum number of times a word must appear in the 
                     corpus to be counted as a feature.
    """
    self._data_fetcher = data_fetcher
    self._cutoff = cutoff
    self._word_count_dict = None


  def gen_features(self, data):
    """Generate features.

    Parameters:
      data   - The data to generate features for.

    Return:
      The feature set for this data.
    """
    if not self._word_count_dict:
      self._create_word_count_dict()
    feature_dict = collections.defaultdict(int)
    for sent in nltk.sent_tokenize(data):
      for word in nltk.word_tokenize(sent):
        word = self._sanitize(word.lower())
        if self._word_count_dict[word] >= self._cutoff and word != '':
          feature_dict[word] += 1
    return feature_dict


  def _create_word_count_dict(self):
    """Calculate the total count for all words in the corpus."""
    self._word_count_dict = collections.defaultdict(int)
    for data, tags in self._data_fetcher.fetch():
      for sent in nltk.sent_tokenize(data):
        for word in nltk.word_tokenize(sent):
          self._word_count_dict[self._sanitize(word.lower())] += 1


################################################################################
# UnigramNegationFeatureGenerator                                              #
# A unigram feature generator that takes into consideration negation words.    #
################################################################################
class UnigramNegationFeatureGenerator(UnigramFeatureGenerator):
  def __init__(self, data_fetcher, cutoff=20):
    """Constructor.

    Parameters:
      data_fetcher - The data_fetcher for the corpus this feature generator will
                     be used with.
      cutoff       - The minimum number of times a word must appear in the 
                     corpus to be counted as a feature.
    """
    super(UnigramNegationFeatureGenerator, self).__init__(data_fetcher, cutoff)
    self._negations = ['not', "n't"]


  def gen_features(self, data):
    """Generate features.

    Parameters:
      data   - The data to generate features for.

    Return:
      The feature set for this data.
    """
    if not self._word_count_dict:
      self._create_word_count_dict()
    feature_dict = collections.defaultdict(int)
    for sent in [nltk.word_tokenize(s) for s in nltk.sent_tokenize(data)]:
      is_negation = self._contains_negation(sent)
      for word in sent:
        word = self._sanitize(word.lower())
        if self._word_count_dict[word] >= self._cutoff and word != '' and \
           not word in self._negations:
            if is_negation:
              word = '<NOT>' + word
            feature_dict[word] += 1
    return feature_dict


  def _contains_negation(self, sent):
    """Check if the sentence contains a negation word.

    Parameters:
      sent - The sentence to check for negation words.

    Return:
      True iff a negation word is present.
    """
    lsent = [w.lower() for w in sent]
    return reduce(operator.or_, [n in lsent for n in self._negations], False)


################################################################################
# AdjectiveFeatureGenerator                                                    #
# A feature generator that only considers adjectives and negation words.       #
################################################################################
class AdjectiveFeatureGenerator(UnigramFeatureGenerator):
  def __init__(self, data_fetcher, cutoff=20):
    """Constructor.

    Parameters:
      data_fetcher - The data_fetcher for the corpus this feature generator will
                     be used with.
      cutoff       - The minimum number of times a word must appear in the 
                     corpus to be counted as a feature.
    """
    super(AdjectiveFeatureGenerator, self).__init__(data_fetcher, cutoff)


  def gen_features(self, data):
    """Generate features.

    Parameters:
      data   - The data to generate features for.

    Return:
      The feature set for this data.
    """
    if not self._word_count_dict:
      self._create_word_count_dict()
    feature_dict = collections.defaultdict(int)
    for sent in [nltk.word_tokenize(s) for s in nltk.sent_tokenize(data)]:
      sent = [(w, simplify_wsj_tag(t)) for (w, t) in nltk.pos_tag(sent)]
      for word, pos in sent:
        if pos == 'ADJ':
          word = self._sanitize(word.lower())
          if self._word_count_dict[word] >= self._cutoff and word != '':
            feature_dict[word] += 1
    return feature_dict


################################################################################
# AdjectiveNegationFeatureGenerator                                            #
# A feature generator that only considers adjectives.                          #
################################################################################
class AdjectiveNegationFeatureGenerator(UnigramNegationFeatureGenerator):
  def __init__(self, data_fetcher, cutoff=20):
    """Constructor.

    Parameters:
      data_fetcher - The data_fetcher for the corpus this feature generator will
                     be used with.
      cutoff       - The minimum number of times a word must appear in the 
                     corpus to be counted as a feature.
    """
    super(AdjectiveNegationFeatureGenerator, self).__init__(
      data_fetcher, cutoff)


  def gen_features(self, data):
    """Generate features.

    Parameters:
      data   - The data to generate features for.

    Return:
      The feature set for this data.
    """
    if not self._word_count_dict:
      self._create_word_count_dict()
    feature_dict = collections.defaultdict(int)
    for sent in [nltk.word_tokenize(s) for s in nltk.sent_tokenize(data)]:
      is_negation = self._contains_negation(sent)
      sent = [(w, simplify_wsj_tag(t)) for (w, t) in nltk.pos_tag(sent)]
      for word, pos in sent:
        if pos == 'ADJ':
          word = self._sanitize(word.lower())
          if self._word_count_dict[word] >= self._cutoff and word != '':
            if is_negation:
              word = '<NOT>' + word
            feature_dict[word] += 1
    return feature_dict
